{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q9KfUf1BI6Kl"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WvqLCVQ6I58i"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZoFRICPTNUca"
      },
      "source": [
        "# Migrate the SavedModel workflow\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/migrate/saved_model\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/saved_model.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/saved_model.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/saved_model.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nGyIb7MYJfaM"
      },
      "source": [
        "Once you have migrated your model from TensorFlow 1's graphs and sessions to TensorFlow 2 APIs, such as `tf.function`, `tf.Module`, and `tf.keras.Model`, you can migrate the model saving and loading code. This notebook provides examples of how you can save and load in the SavedModel format in TensorFlow 1 and TensorFlow 2. Here is a quick overview of the related API changes for migration from TensorFlow 1 to TensorFlow 2:\n",
        "\n",
        "|                                                                                         | TensorFlow 1                                                                                                                                                         | Migration to TensorFlow 2                                                                                                                                                                  |\n",
        "| --------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |\n",
        "| **Saving**                                                                              | `tf.compat.v1.saved_model.Builder`<br>`tf.compat.v1.saved_model.simple_save`                                                                                         | `tf.saved_model.save`<br>Keras: `tf.keras.models.save_model`                                                                                                                               |\n",
        "| **Loading**                                                                             | `tf.compat.v1.saved_model.load`                                                                                                                                      | `tf.saved_model.load`<br>Keras: `tf.keras.models.load_model`                                                                                                                               |\n",
        "| **Signatures**: a set of input<br>and output tensors that<br>can be used to run the<br> | Generated using the `*.signature_def` utils<br>(e.g. `tf.compat.v1.saved_model.predict_signature_def`)                                                               | Write a `tf.function` and export it using the `signatures` argument<br>in `tf.saved_model.save`.                                                                                           |\n",
        "| **Classification<br>and regression**:<br>special types of signatures                    | Generated with<br>`tf.compat.v1.saved_model.classification_signature_def`,<br>`tf.compat.v1.saved_model.regression_signature_def`,<br>and certain Estimator exports. | These two signature types have been removed from TensorFlow 2.<br>If the serving library requires these method names,<br>`tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater`. |\n",
        "\n",
        "For a more in-depth explanation of the mapping, refer to the [Changes from TensorFlow 1 to TensorFlow 2](#changes_from_tf1_to_tf2) section below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r5mR2xsNAGsB"
      },
      "source": [
        "## Setup\n",
        "\n",
        "The examples below show how to export and load the same dummy TensorFlow model (defined as `add_two` below) to a SavedModel format using the TensorFlow 1 and TensorFlow 2 APIs. Start by setting up the imports and utility functions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B94QZyy-kOGQ"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1\n",
        "import shutil\n",
        "\n",
        "def remove_dir(path):\n",
        "  try:\n",
        "    shutil.rmtree(path)\n",
        "  except:\n",
        "    pass\n",
        "\n",
        "def add_two(input):\n",
        "  return input + 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZNVpH5tOCgd9"
      },
      "source": [
        "## TensorFlow 1: Save and export a SavedModel\n",
        "\n",
        "In TensorFlow 1, you use the `tf.compat.v1.saved_model.Builder`, `tf.compat.v1.saved_model.simple_save`, and `tf.estimator.Estimator.export_saved_model` APIs to build, save, and export the TensorFlow graph and session:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "THRLul5ijmTE"
      },
      "source": [
        "### 1. Save the graph as a SavedModel with SavedModelBuilder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dcZDQaI8jl3h"
      },
      "outputs": [],
      "source": [
        "remove_dir(\"saved-model-builder\")\n",
        "\n",
        "with tf.Graph().as_default() as g:\n",
        "  with tf1.Session() as sess:\n",
        "    input = tf1.placeholder(tf.float32, shape=[])\n",
        "    output = add_two(input)\n",
        "    print(\"add two output: \", sess.run(output, {input: 3.}))\n",
        "\n",
        "    # Save with SavedModelBuilder\n",
        "    builder = tf1.saved_model.Builder('saved-model-builder')\n",
        "    sig_def = tf1.saved_model.predict_signature_def(\n",
        "        inputs={'input': input},\n",
        "        outputs={'output': output})\n",
        "    builder.add_meta_graph_and_variables(\n",
        "        sess, tags=[\"serve\"], signature_def_map={\n",
        "            tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: sig_def\n",
        "    })\n",
        "    builder.save()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PwtC27VFlwCa"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir saved-model-builder --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gnBDNTxKG_vR"
      },
      "source": [
        "### 2. Build a SavedModel for serving"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jtMxe2rjHSq9"
      },
      "outputs": [],
      "source": [
        "remove_dir(\"simple-save\")\n",
        "\n",
        "with tf.Graph().as_default() as g:\n",
        "  with tf1.Session() as sess:\n",
        "    input = tf1.placeholder(tf.float32, shape=[])\n",
        "    output = add_two(input)\n",
        "    print(\"add_two output: \", sess.run(output, {input: 3.}))\n",
        "\n",
        "    tf1.saved_model.simple_save(\n",
        "        sess, 'simple-save',\n",
        "        inputs={'input': input},\n",
        "        outputs={'output': output})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AdnqemvIHb2P"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir simple-save --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r0BNzzAHjnkp"
      },
      "source": [
        "### 3. Export the Estimator inference graph as a SavedModel\n",
        "\n",
        "In the definition of the Estimator `model_fn` (defined below), you can define signatures in your model by returning `export_outputs` in the `tf.estimator.EstimatorSpec`. There are different types of outputs:\n",
        "\n",
        "- `tf.estimator.export.ClassificationOutput`\n",
        "- `tf.estimator.export.RegressionOutput`\n",
        "- `tf.estimator.export.PredictOutput`\n",
        "\n",
        "These will produce classification, regression, and prediction signature types, respectively.\n",
        "\n",
        "When the estimator is exported with `tf.estimator.Estimator.export_saved_model`, these signatures will be saved with the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3nQ5Stnxjhfs"
      },
      "outputs": [],
      "source": [
        "def model_fn(features, labels, mode):\n",
        "  output = add_two(features['input'])\n",
        "  step = tf1.train.get_global_step()\n",
        "  return tf.estimator.EstimatorSpec(\n",
        "      mode,\n",
        "      predictions=output,\n",
        "      train_op=step.assign_add(1),\n",
        "      loss=tf.constant(0.),\n",
        "      export_outputs={\n",
        "          tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: \\\n",
        "          tf.estimator.export.PredictOutput({'output': output})})\n",
        "est = tf.estimator.Estimator(model_fn, 'estimator-checkpoints')\n",
        "\n",
        "# Train for one step to create a checkpoint.\n",
        "def train_fn():\n",
        "  return tf.data.Dataset.from_tensors({'input': 3.})\n",
        "est.train(train_fn, steps=1)\n",
        "\n",
        "# This utility function `build_raw_serving_input_receiver_fn` takes in raw\n",
        "# tensor features and builds an \"input serving receiver function\", which\n",
        "# creates placeholder inputs to the model.\n",
        "serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(\n",
        "    {'input': tf.constant(3.)})  # Pass in a dummy input batch.\n",
        "estimator_path = est.export_saved_model('exported-estimator', serving_input_fn)\n",
        "\n",
        "# Estimator's export_saved_model creates a time stamped directory. Move this\n",
        "# to a set path so it can be inspected with `saved_model_cli` in the cell below.\n",
        "!rm -rf estimator-model\n",
        "import shutil\n",
        "shutil.move(estimator_path, 'estimator-model')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8_gD2gkE7CMu"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir estimator-model --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=[10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DyBvrNQgIhIo"
      },
      "source": [
        "## TensorFlow 2: Save and export a SavedModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZmFH-eIjqjB"
      },
      "source": [
        "### Save and export a SavedModel defined with tf.Module\n",
        "\n",
        "To export your model in TensorFlow 2, you must define a `tf.Module` or a `tf.keras.Model` to hold all of your model's variables and functions. Then, you can call `tf.saved_model.save` to create a SavedModel. Refer to the _Saving a custom model_ section in the [Using the SavedModel format](../saved_model.ipynb) guide to learn more."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_j-PwgP_jrgw"
      },
      "outputs": [],
      "source": [
        "class MyModel(tf.Module):\n",
        "  @tf.function\n",
        "  def __call__(self, input):\n",
        "    return add_two(input)\n",
        "\n",
        "model = MyModel()\n",
        "\n",
        "@tf.function\n",
        "def serving_default(input):\n",
        "  return {'output': model(input)}\n",
        "\n",
        "signature_function = serving_default.get_concrete_function(\n",
        "    tf.TensorSpec(shape=[], dtype=tf.float32))\n",
        "tf.saved_model.save(\n",
        "    model, 'tf2-save', signatures={\n",
        "        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_function})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "slvU4vZN756F"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir tf2-save --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYpSfbBJjr33"
      },
      "source": [
        "### Save and export a SavedModel defined with Keras\n",
        "\n",
        "The Keras APIs for saving and exporting—`Mode.save` or `tf.keras.models.save_model`—can export a SavedModel from a `tf.keras.Model`. Check out the [Save and load Keras models](../..guide/keras/save_and_serialize) for more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mMcjhzyRjvp6"
      },
      "outputs": [],
      "source": [
        "inp = tf.keras.Input(3)\n",
        "out = add_two(inp)\n",
        "model = tf.keras.Model(inputs=inp, outputs=out)\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])\n",
        "def serving_default(input):\n",
        "  return {'output': model(input)}\n",
        "\n",
        "model.save('keras-model', save_format='tf', signatures={\n",
        "        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: serving_default})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4P93WP5R7-VT"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir keras-model --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SEKe9rGgoGCw"
      },
      "source": [
        "## Loading a SavedModel\n",
        "\n",
        "A SavedModel saved with any of the above APIs can be loaded using either TensorFlow 1 or TensorFlow 2 APIs.\n",
        "\n",
        "A TensorFlow 1 SavedModel can generally be used for inference when loaded into TensorFlow 2, but training (generating gradients) is only possible if the SavedModel contains *resource variables*. You can check the dtype of the variables—if the variable dtype contains \"_ref\", then it is a reference variable.\n",
        "\n",
        "A TensorFlow 2 SavedModel can be loaded and executed from TensorFlow 1 as long as the SavedModel is saved with signatures.\n",
        "\n",
        "The sections below contain code samples showing how to load the SavedModels saved in the previous sections, and call the exported signature."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hLztK_0YoTEP"
      },
      "source": [
        "### TensorFlow 1: Load a SavedModel with tf.saved_model.load\n",
        "\n",
        "In TensorFlow 1, you can import a SavedModel directly into the current graph and session using `tf.saved_model.load`. You can call `Session.run` on the tensor input and output names:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IMO0laj-m0p9"
      },
      "outputs": [],
      "source": [
        "def load_tf1(path, input):\n",
        "  print('Loading from', path)\n",
        "  with tf.Graph().as_default() as g:\n",
        "    with tf1.Session() as sess:\n",
        "      meta_graph = tf1.saved_model.load(sess, [\"serve\"], path)\n",
        "      sig_def = meta_graph.signature_def[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]\n",
        "      input_name = sig_def.inputs['input'].name\n",
        "      output_name = sig_def.outputs['output'].name\n",
        "      print('  Output with input', input, ': ', \n",
        "            sess.run(output_name, feed_dict={input_name: input}))\n",
        "\n",
        "load_tf1('saved-model-builder', 5.)\n",
        "load_tf1('simple-save', 5.)\n",
        "load_tf1('estimator-model', [5.])  # Estimator's input must be batched.\n",
        "load_tf1('tf2-save', 5.)\n",
        "load_tf1('keras-model', 5.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FbR3sfvooVBN"
      },
      "source": [
        "### TensorFlow 2: Load a model saved with tf.saved_model\n",
        "\n",
        "In TensorFlow 2, objects are loaded into a Python object that stores the variables and functions. This is compatible with models saved from TensorFlow 1.\n",
        "\n",
        "Check out the `tf.saved_model.load` API docs and [Loading and using a custom model](../../guide/saved_model#loading_and_using_a_custom_model) section from the [Using the SavedModel format](../..guide/saved_model) guide for details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OA52ezWV_KgL"
      },
      "outputs": [],
      "source": [
        "def load_tf2(path, input):\n",
        "  print('Loading from', path)\n",
        "  loaded = tf.saved_model.load(path)\n",
        "  out = loaded.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY](\n",
        "      tf.constant(input))['output']\n",
        "  print('  Output with input', input, ': ', out)\n",
        "\n",
        "load_tf2('saved-model-builder', 5.)\n",
        "load_tf2('simple-save', 5.)\n",
        "load_tf2('estimator-model', [5.])  # Estimator's input must be batched.\n",
        "load_tf2('tf2-save', 5.)\n",
        "load_tf2('keras-model', 5.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gz3VFn5aAfmK"
      },
      "source": [
        "Models saved with the TensorFlow 2 API can also access `tf.function`s and variables that are attached to the model (instead of those exported as signatures). For example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IfMTp-TGAfOs"
      },
      "outputs": [],
      "source": [
        "loaded = tf.saved_model.load('tf2-save')\n",
        "print('restored __call__:', loaded.__call__)\n",
        "print('output with input 5.', loaded(5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VMoErNKHoXEg"
      },
      "source": [
        "### TensorFlow 2: Load a model saved with Keras\n",
        "\n",
        "The Keras loading API—`tf.keras.models.load_model`—allows you to reload a saved model back into a Keras Model object. Note that this only allows you to load SavedModels saved with Keras (`Model.save` or `tf.keras.models.save_model`).\n",
        "\n",
        "Models saved with `tf.saved_model.save` should be loaded with `tf.saved_model.load`. You can load a Keras model saved with `Model.save` using `tf.saved_model.load` but you will only get the TensorFlow graph. Refer to the `tf.keras.models.load_model` API docs and [Save and load Keras models](https://www.tensorflow.org/guide/keras/save_and_serialize#savedmodel_format) guide for details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZFUAxK0YeIAe"
      },
      "outputs": [],
      "source": [
        "loaded_model = tf.keras.models.load_model('keras-model')\n",
        "loaded_model.predict_on_batch(tf.constant([1, 3, 4]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tz4eAAGY19MM"
      },
      "source": [
        "## GraphDef and MetaGraphDef\n",
        "\n",
        "<a name=\\\"graphdef_and_metagraphdef\\\"> </a>\n",
        "\n",
        "There is no straightforward way to load a raw `GraphDef` or `MetaGraphDef` to TF2. However, you can convert the TF1 code that imports the graph into a TF2 [`concrete_function`](https://tensorflow.org/guide/concrete_function) using [`v1.wrap_function`](https://www.tensorflow.org/api_docs/python/tf/compat/v1/wrap_function).\n",
        "\n",
        "First, save a MetaGraphDef:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "grKam9zGnNRZ"
      },
      "outputs": [],
      "source": [
        "# Save a simple multiplication computation:\n",
        "with tf.Graph().as_default() as g:\n",
        "  x = tf1.placeholder(tf.float32, shape=[], name='x')\n",
        "  v = tf.Variable(3.0, name='v')\n",
        "  y = tf.multiply(x, v, name='y')\n",
        "  with tf1.Session() as sess:\n",
        "    sess.run(v.initializer)\n",
        "    print(sess.run(y, feed_dict={x: 5}))\n",
        "    s = tf1.train.Saver()\n",
        "    s.export_meta_graph('multiply.pb', as_text=True)\n",
        "    s.save(sess, 'multiply_values.ckpt')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bJnCe7eYrXev"
      },
      "source": [
        "Using TF1 APIs, you can use `tf1.train.import_meta_graph` to import the graph and restore the values:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3bbcGZ4CoVDL"
      },
      "outputs": [],
      "source": [
        "with tf.Graph().as_default() as g:\n",
        "  meta = tf1.train.import_meta_graph('multiply.pb')\n",
        "  x = g.get_tensor_by_name('x:0')\n",
        "  y = g.get_tensor_by_name('y:0')\n",
        "  with tf1.Session() as sess:\n",
        "    meta.restore(sess, 'multiply_values.ckpt')\n",
        "    print(sess.run(y, feed_dict={x: 5}))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s9OnHOLDrnco"
      },
      "source": [
        "There are no TF2 APIs for loading the graph, but you can still import it into a concrete function that can be executed in eager mode:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yAyGmDLlpVBX"
      },
      "outputs": [],
      "source": [
        "def import_multiply():\n",
        "  # Any graph-building code is allowed here.\n",
        "  tf1.train.import_meta_graph('multiply.pb')\n",
        "\n",
        "# Creates a tf.function with all the imported elements in the function graph.\n",
        "wrapped_import = tf1.wrap_function(import_multiply, [])\n",
        "import_graph = wrapped_import.graph\n",
        "x = import_graph.get_tensor_by_name('x:0')\n",
        "y = import_graph.get_tensor_by_name('y:0')\n",
        "\n",
        "# Restore the variable values.\n",
        "tf1.train.Saver(wrapped_import.variables).restore(\n",
        "    sess=None, save_path='multiply_values.ckpt')\n",
        "\n",
        "# Create a concrete function by pruning the wrap_function (similar to sess.run).\n",
        "multiply_fn = wrapped_import.prune(feeds=x, fetches=y)\n",
        "\n",
        "# Run this function\n",
        "multiply_fn(tf.constant(5.))  # inputs to concrete functions must be Tensors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZ5vGJ0IDorc"
      },
      "source": [
        "## Changes from TensorFlow 1 to TensorFlow 2\n",
        "\n",
        "<a id=\\\"changes_from_tf1_to_tf2\\\"> </a>\n",
        "\n",
        "This section lists out key saving and loading terms from TensorFlow 1, their TensorFlow 2 equivalents, and what has changed.\n",
        "\n",
        "### SavedModel\n",
        "\n",
        "[SavedModel](../../guide/saved_model.ipynb) is a format that stores a complete TensorFlow program with parameters and computation. It contains signatures used by serving platforms to run the model.\n",
        "\n",
        "The file format itself has not changed significantly, so SavedModels can be loaded and served using either TensorFlow 1 or TensorFlow 2 APIs.\n",
        "\n",
        "**Differences between TensorFlow 1 and TensorFlow 2**\n",
        "\n",
        "The *serving* and *inference* use cases have not been updated in TensorFlow 2, aside from the API changes—the improvement was introduced in the ability to *reuse* and *compose models* loaded from SavedModel.\n",
        "\n",
        "In TensorFlow 2, the program is represented by objects like `tf.Variable`, `tf.Module`, or higher-level Keras models (`tf.keras.Model`) and layers (`tf.keras.layers`). There are no more global variables that have values stored in a session, and the graph now exists in different `tf.function`s. Consequently, during a model export, SavedModel saves each component and function graphs separately.\n",
        "\n",
        "When you write a TensorFlow program with the TensorFlow Python APIs, you must build an object to manage the variables, functions, and other resources. Generally, this is accomplished by using the Keras APIs, but you can also build the object by creating or subclassing `tf.Module`.\n",
        "\n",
        "Keras models (`tf.keras.Model`) and `tf.Module` automatically track variables and functions attached to them. SavedModel saves these connections between modules, variables, and functions, so that they can be restored when loading.\n",
        "\n",
        "### Signatures\n",
        "\n",
        "Signatures are the endpoints of a SavedModel—they tell the user how to run the model and what inputs are needed.\n",
        "\n",
        "In TensorFlow 1, signatures are created by listing the input and output tensors. In TensorFlow 2, signatures are generated by passing in *concrete functions*. (Read more about TensorFlow functions in the [Introduction to graphs and tf.function](../intro_to_graphs.ipynb) guide, particularly the _Polymorphism: one Function, many graphs_ section.) In short, a concrete function is generated from a `tf.function`:\n",
        "\n",
        "```python\n",
        "# Option 1: Specify an input signature.\n",
        "@tf.function(input_signature=[...])\n",
        "def fn(...):\n",
        "  ...\n",
        "  return outputs\n",
        "\n",
        "tf.saved_model.save(model, path, signatures={\n",
        "    'name': fn\n",
        "})\n",
        "```\n",
        "\n",
        "```python\n",
        "# Option 2: Call `get_concrete_function`\n",
        "@tf.function\n",
        "def fn(...):\n",
        "  ...\n",
        "  return outputs\n",
        "\n",
        "tf.saved_model.save(model, path, signatures={\n",
        "    'name': fn.get_concrete_function(...)\n",
        "})\n",
        "```\n",
        "\n",
        "### `Session.run`\n",
        "\n",
        "In TensorFlow 1, you could call `Session.run` with the imported graph as long as you already know the tensor names. This allows you to retrieve the restored variable values, or run parts of the model that were not exported in the signatures.\n",
        "\n",
        "In TensorFlow 2, you can directly access a variable, such as a weights matrix (`kernel`):\n",
        "\n",
        "```python\n",
        "model = tf.Module()\n",
        "model.dense_layer = tf.keras.layers.Dense(...)\n",
        "tf.saved_model.save('my_saved_model')\n",
        "loaded = tf.saved_model.load('my_saved_model')\n",
        "loaded.dense_layer.kernel\n",
        "```\n",
        "\n",
        "\n",
        "or call `tf.function`s attached to the model object: for example, `loaded.__call__`.\n",
        "\n",
        "Unlike TF1, there is no way to extract parts of a function and access intermediate values. You *must* export all of the needed functionality in the saved object.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b6NG9JvUwJxn"
      },
      "source": [
        "## TensorFlow Serving migration notes\n",
        "\n",
        "SavedModel was originally created to work with [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving). This platform offers different types of prediction requests: classify, regress, and predict.\n",
        "\n",
        "The **TensorFlow 1** API allows you to create these types of signatures with the utils:\n",
        "\n",
        "- `tf.compat.v1.saved_model.classification_signature_def`\n",
        "- `tf.compat.v1.saved_model.regression_signature_def`\n",
        "- `tf.compat.v1.saved_model.predict_signature_def`\n",
        "\n",
        "[Classification](https://www.tensorflow.org/tfx/serving/signature_defs#classification_signaturedef) (`classification_signature_def`) and [regression](https://www.tensorflow.org/tfx/serving/signature_defs#regression_signaturedef) (`regression_signature_def`) restrict the inputs and outputs, so the inputs must be a `tf.Example`, and the outputs must be `classes`, `scores` or `prediction`. Meanwhile, [the predict signature](https://www.tensorflow.org/tfx/serving/signature_defs#predict_signaturedef) (`predict_signature_def`) has no restrictions.\n",
        "\n",
        "SavedModels exported with the **TensorFlow 2** API are compatible with TensorFlow Serving, but will only contain prediction signatures. The classification and regression signatures have been removed.\n",
        "\n",
        "If you require the use of the classification and regression signatures, you may modify the exported SavedModel using `tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a3acd3b86215"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "To learn more about SavedModels in TensorFlow 2, check out the following guides:\n",
        "\n",
        "- [Using the SavedModel format](https://www.tensorflow.org/guide/saved_model)\n",
        "- [Save and load Keras models](https://www.tensorflow.org/guide/keras/save_and_serialize)\n",
        "\n",
        "If you are using TensorFlow Hub, you may find these guides useful:\n",
        "\n",
        "- [TensorFlow Hub: Model compatibility for TensorFlow 1/TensorFlow 2](https://www.tensorflow.org/hub/model_compatibility)\n",
        "- [Migrating from TensorFlow 1 to TensorFlow 2 with TensorFlow Hub](https://www.tensorflow.org/hub/migration_tf2)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "saved_model.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
